# coding=utf-8
"""
    @project: MaxKB
    @Author：虎
    @file： mark_chunk_handle.py
    @date：2024/7/23 16:52
    @desc:
"""
import re
from typing import List

from common.chunk.i_chunk_handle import IChunkHandle

max_chunk_len = 128
split_chunk_pattern = r'.{1,%d}[。| |\\.|！|;|；|!|\n]' % max_chunk_len
max_chunk_pattern = r'.{1,%d}' % max_chunk_len


class MarkChunkHandle(IChunkHandle):
    def handle(self, chunk_list: List[str], chunk_patterns: str = None, chunk_length: int = None):
        if chunk_length is None:
            chunk_length = max_chunk_len
        if chunk_patterns is None:
            chunk_patterns = r'.{1,%d}[。| |\\.|！|;|；|!|\n]' % chunk_length
        else:
            # 转义用户输入的特殊字符
            escaped_patterns = re.escape(chunk_patterns)
            # 构建新的正则表达式模式
            chunk_patterns = r'.{1,%d}[%s]' % (chunk_length, escaped_patterns)
            
        mcp = r'.{1,%d}' % chunk_length
            
        result = []
        for chunk in chunk_list:
            chunk_result = re.findall(chunk_patterns, chunk, flags=re.DOTALL)
            for c_r in chunk_result:
                if len(c_r.strip()) > 0:
                    result.append(c_r.strip())

            other_chunk_list = re.split(chunk_patterns, chunk, flags=re.DOTALL)
            for other_chunk in other_chunk_list:
                if len(other_chunk) > 0:
                    if len(other_chunk) < chunk_length:
                        if len(other_chunk.strip()) > 0:
                            result.append(other_chunk.strip())
                    else:
                        max_chunk_list = re.findall(mcp, other_chunk, flags=re.DOTALL)
                        for m_c in max_chunk_list:
                            if len(m_c.strip()) > 0:
                                result.append(m_c.strip())

        return result
